# -*- coding: utf-8 -*-
"""
Created on Tue Jun 12 09:36:55 2018
func：加载模型，进行模型测试
@author: kuangyongjian
"""
import tensorflow as tf
import numpy as np
from PIL import Image
from model import Network
import os.path
import re
import matplotlib.pyplot as plt
import matplotlib.pylab as pylab
CKPT_DIR = 'ckpt'

# turn picture to 28*28
def img_resize(img_file, path_save, width,height):
    img = Image.open(img_file)
    new_image = img.resize((width,height),Image.BILINEAR)
    new_image.save(os.path.join(path_save,os.path.basename(img_file)))
 
class Predict(object):
    
    def __init__(self):
        #清除默认图的堆栈，并设置全局图为默认图
        #若不进行清楚则在第二次加载的时候报错，因为相当于重新加载了两次
        tf.reset_default_graph() 
        self.net = Network()
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())
        
        #加载模型到sess中
        self.restore()
        print('load susess')
    
    def restore(self):
        saver = tf.train.Saver()
        ckpt = tf.train.get_checkpoint_state(CKPT_DIR)
        print(ckpt.model_checkpoint_path)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(self.sess,ckpt.model_checkpoint_path)
        else:
            raise FileNotFoundError('未保存模型')
        
    def predict(self,image_path):
        #读取图片并灰度化
        img_resize(image_path, './input_picture',28,28) 
        get_name = re.compile(r'[^/]+\.png',re.M)
        picture_name = get_name.search(image_path).group()
        path = './input_picture/' + picture_name
        img = Image.open(path).convert('L')
        flatten_img = np.reshape(img,784)
        x = np.array([1 - flatten_img])
        y = self.sess.run(self.net.y,feed_dict = {self.net.x:x})
        

        # pic_matrix = np.matrix(one_pic_arr, dtype="float")
        plt.imshow(img)
        pylab.show()
        print(image_path)
        print(' Predict digit:  ',np.argmax(y[0]))
        
        
if __name__ == '__main__':
    model = Predict()
    model.predict('./source_picture/1.png')
    model.predict('./source_picture/2.png')
    model.predict('./source_picture/3.png')

  # # show the gragh
  # for i in range(0, len(mnist.test.images)):
  #   result = sess.run(correct_prediction, feed_dict={x: np.array([mnist.test.images[i]]), y_: np.array([mnist.test.labels[i]])})
  #   if not result:
  #     print('预测的值是：',sess.run(y, feed_dict={x: np.array([mnist.test.images[i]]), y_: np.array([mnist.test.labels[i]])}))
  #     print('实际的值是：',sess.run(y_,feed_dict={x: np.array([mnist.test.images[i]]), y_: np.array([mnist.test.labels[i]])}))
  #     one_pic_arr = np.reshape(mnist.test.images[i], (28, 28))
  #     pic_matrix = np.matrix(one_pic_arr, dtype="float")
  #     plt.imshow(pic_matrix)
  #     pylab.show()
  #     break